package uk.ac.aber.paws.server.security;

import java.io.BufferedInputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.FileWriter;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.net.InetAddress;
import java.net.URL;
import java.net.URLConnection;
import java.util.Vector;
import java.security.*;
import java.security.cert.*;
import javax.net.ssl.*;

import uk.ac.aber.paws.server.config.ConfigManager;
import uk.ac.aber.paws.server.core.PAWSServer;

/**
 * @author Rob Johnson, Aberystwyth University
 * @email rob.johnson@aber.ac.uk
 * @date 21st May 2012
 * 
 * SSLInitialise.java
 * 
 * Initialises various aspects of the trusted SSL certs witin the application
 * 
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
public class SSLInitialise {
 
	private static Vector<Boolean> sslStatus = new Vector<Boolean>();
	private static Vector<X509Certificate> certList = new Vector<X509Certificate>();
	
	/**
	 * Creates a Java keystore in the users home directory
	 */
	public static void createKeystore() {
		
		String fileName = System.getProperty("user.home") +
		  System.getProperty("file.separator") + "paws.keystore";
		String msg  = "This message was written by a signed applet!!!\n";
		String s ;
				  
		try {
		FileWriter fos = new FileWriter(fileName);
		fos.write(msg, 0, msg.length());
		fos.close();
		s = new String("Successfully created file :" + fileName);
		
		} catch (Exception e) {
		System.out.println("Exception e = " + e);
		e.printStackTrace();
		s = new String("Unable to create file :  " + fileName);
		}
	}
	
	/**
	 * Retrieves SSL cert from the specified domain controller and imports it into the users keystore
	 * @param ldapServer The domain controller fqdn
	 * @param ldapPort The port on which the LDAPS server is runnign
	 * @param keystorePass The password to the keystore
	 */
	public static void retrieveLDAPServerCerts(String ldapServer, int ldapPort, String keystorePass)  {
		
		Vector<String> serverList = dnsLookup(ldapServer);
		
		//createKeystore();
		
		if (serverList != null) {
			
			for (int i=0;i<serverList.size();i++) {
				
				try {
					checkCert(serverList.get(i),ldapPort,keystorePass);
					sslStatus.add(true);
				}
				catch (Exception ex) {
					ex.printStackTrace();
					sslStatus.add(false);
				}
			}
		}
	}
	
	/**
	 * Returns a list of the stored SSL certificates
	 * @return
	 */
	public static Vector<X509Certificate> getCertList() {
	
		return certList;
	}
	
	/**
	 * Returns true is all SSL certs were imported successfully and false if any failed
	 */
	public static boolean didSSLFail() {
	
		for (int i=0;i<sslStatus.size();i++) {
			
			if (!sslStatus.get(i)) {
				return true;
			}
		}
		
		return false;
	}
	
	/**
	 * Downloads a template paws java keystore file
	 * @param url The URL to the hosted template file via HTTP
	 * @param localFile The local file in which to write the JKS file
	 */
	private static void downloadCert(URL url, String localFile) {
		
		try {
		    URLConnection uc = url.openConnection();
		   
		    int contentLength = uc.getContentLength();
		    
		    InputStream inputStream = uc.getInputStream();
		    InputStream in = new BufferedInputStream(inputStream);
		    
		    byte[] data = new byte[contentLength];
		    int bytesRead = 0;
		    int offset = 0;
		    
		    while (offset < contentLength) {
		    	
		    	bytesRead = in.read(data, offset, data.length - offset);
		    	
		    	if (bytesRead == -1) {
		    		break;
		    	}
		    	
		    	offset += bytesRead;
		    }
		    
		    in.close();

		    File updateFile = new File(localFile);
		    
		    updateFile.mkdirs();
		    
		    if (updateFile.exists()) {
		    	updateFile.delete();
		    }
		    
		    FileOutputStream out = new FileOutputStream(updateFile);
		    
		    out.write(data);
		    out.flush();
		    out.close();
		}
		catch (IOException ex) {
			ex.printStackTrace();
		}
	}

	/**
	 * Downloads from the DCs, checks and imports the SSL certs into the keystore
	 * @param host The host in which to download the SSL cert from
	 * @param checkPort The LDAPS port in which to bind to for the cert
	 * @param password The keystore password, required to import the SSL certs
	 * @throws Exception
	 */
	public static void checkCert(String host, int checkPort, String password) throws Exception {
		    
		password = "kMdCVb5";
		
    	char[] passphrase = password.toCharArray();
	
    	File file = null;
   
    	file = new File(System.getProperty("user.home")+"/paws.jks");
    	
    	System.setProperty("javax.net.ssl.trustStore", System.getProperty("user.home")+"/paws.jks");
    	
    	if (!file.exists()) {
    		downloadCert(new URL(ConfigManager.getProperty("app.ssl.remotejksfile")),System.getProperty("user.home")+"/paws.jks");
    	}
    	
		InputStream in = new FileInputStream(file);
		KeyStore ks = KeyStore.getInstance(KeyStore.getDefaultType());
		ks.load(in, passphrase);
		in.close();

		
		SSLContext context = SSLContext.getInstance("TLS");
		TrustManagerFactory tmf = TrustManagerFactory.getInstance(TrustManagerFactory.getDefaultAlgorithm());
		tmf.init(ks);
		
		X509TrustManager defaultTrustManager = (X509TrustManager)tmf.getTrustManagers()[0];
		SavingTrustManager tm = new SavingTrustManager(defaultTrustManager);
		context.init(null, new TrustManager[] {tm}, null);
		SSLSocketFactory factory = context.getSocketFactory();

		SSLSocket socket = (SSLSocket)factory.createSocket(host, checkPort);
		socket.setSoTimeout(10000);

		boolean check = false;
		
		try {
		    socket.startHandshake();
		    socket.close();		    
		    check = true;
		
		    X509Certificate[] tmpArray = tm.getChain();
		    
		    System.out.println("Current certificate list: "+certList.toString());
		    
		    for (int i=0;i<tmpArray.length;i++) {
		    	certList.add((tmpArray[i]));
			}
			
		    System.out.println("Current certificate list now: "+certList.toString());
			   
		} 
		catch (SSLException e) {
			e.printStackTrace();
		    System.out.println("Could not find certificate for: "+host);
		}
   
		if (!check) {
			X509Certificate[] chain = tm.getChain();
			
			if (chain == null) {
			    System.out.println("Could not obtain server certificate chain");
			   
			}
		
			MessageDigest sha1 = MessageDigest.getInstance("SHA1");
			MessageDigest md5 = MessageDigest.getInstance("MD5");
			for (int i = 0; i < chain.length; i++) {
			    X509Certificate cert = chain[i];
			    sha1.update(cert.getEncoded());
			    md5.update(cert.getEncoded());
			 
			}
		
			int k = 0;
			
			X509Certificate cert = chain[k];
			String alias = host + "-" + (k + 1);
			ks.setCertificateEntry(alias, cert);
		
			OutputStream out = new FileOutputStream(file);
			ks.store(out, passphrase);
			out.close();
		
			certList.add(cert);
			System.out.println("Added certificate to keystore cacerts using alias '"+ alias + "' for host "+host);
		}
    }
	 
	/**
	 * Returns a list of IP addresses from a DNS lookup of the host (could be DNS round robin or singular server)
	 * @param fqdn The FQDN of the server or round robin
	 * @return
	 */
	private static Vector<String> dnsLookup(String fqdn) {
		
	    Vector<String> serverList = new Vector<String>();
	    
	    try {
	    	InetAddress[] computerList = InetAddress.getAllByName(fqdn);
	    	
	    	for (int i=0;i<computerList.length;i++) {
	    		InetAddress ip = computerList[i];
	    		
	    		serverList.add(ip.getHostAddress());
	    	}
	    } 
	    catch (Exception ue) {
	    	ue.printStackTrace();
	    	System.out.println("Cannot find host " + fqdn);
	    	return null;
	    }
	    
	    return serverList;
	}
		
}
